import torchvision.models as models
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from transformers import CLIPProcessor, CLIPModel         # 导入多模态模型

from dataloader import get_dataset, get_dataloader
from Resnet18 import Classifier, Combied_res_clas, Combied_res_CLIP, evaluate, get_largest_pth_file, save_txt
import matplotlib.pyplot as plt
import time  
from tqdm import tqdm
from torch.optim import lr_scheduler     # 调整学习率


torch.backends.cudnn.benchmark = True           # 优化内存分配方式
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 避免死锁  配置线程


start_t1 = time.time()   # 记录开始时间

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#  加载网络模型 resnet18
resnet18 = models.resnet18(weights=None)
resnet18.fc = nn.Linear(resnet18.fc.in_features, 512, bias=False)
resnet18 = resnet18.to(device)  # 转换模型到 GPU
# print("resnet完整结构: ", resnet18)

# 加载多模态模型
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
# 不更新模型的参数
for param in model.parameters():
    param.requires_grad = False

start_t2 = time.time()    # 记录加载模型的时间
print("加载模型的时间：", (start_t2-start_t1)/60)


# 加载数据集
cifar100_dataset = get_dataset() 
cifar100_dataset_test = get_dataset(is_train=False) 
cifar100_classes_name, cifar100_dataloader = get_dataloader(cifar100_dataset)
_, cifar100_dataloader_test = get_dataloader(cifar100_dataset_test)

# 加载模型
classifiar = Classifier(100)
combied_res_clas = Combied_res_clas(resnet18, classifiar).to(device)
combied_res_clip = Combied_res_CLIP(resnet18).to(device)
# 如果存在则加载 
if os.path.isdir(".mdeol/save/cifar100"):
    path = get_largest_pth_file(".mdeol/save/cifar100")
    combied_res_clip.load_state_dict(torch.load(path))

start_t3 = time.time()   # 记录加载数据集的时间
print("加载数据集的时间: ", (start_t3-start_t2)/60)

# 配置优化器、损失函数
cross_loss = nn.CrossEntropyLoss()
cross_optimizer = optim.Adam(combied_res_clas.parameters(), lr=0.001)     # resnet
clip_optimizer = optim.Adam(combied_res_clip.parameters(), lr=0.001)      # clip
# 对学习率进行调整
cross_scheduler = lr_scheduler.StepLR(cross_optimizer, step_size=50, gamma=0.5)  # 每50个epoch 学习率减半

train_clip_losses =   []      # 记录clip损失
trian_resnet_losses = []      # 记录resnet损失
train_accuracies = []         # 记录 训练准确率
test_accuracies = []          # 记录 测试准确率

num_epochs = 200


# 开始训练
def train(epoch):
    combied_res_clas.train()
    # 分别记录 clip损失 和 resnet 损失  正确数量
    total_clip_loss, total_resnet_loss, correct = 0, 0, 0 

    for orig_input, label in tqdm(cifar100_dataloader, desc=f"Epoch {epoch}", unit="batch"):
        text = [f"a photo of the {cifar100_classes_name[i]}"for i in label]
        # print(text)
        orig_input, label = orig_input.to(device), label.to(device)
        clip_inputs = processor(text=text, images=orig_input, return_tensors="pt", padding=True, do_rescale=False)
        # 将clip_inputs中的每个张量都移动到GPU设备
        clip_inputs = {k: v.to(device) for k, v in clip_inputs.items()}
        clip_outputs = model(**clip_inputs)
        clip_img_embed = clip_outputs.image_embeds   # 获取图像编码
        # resnet_embed = resnet18(orig_input) 


        # 组合损失 反向传播
        clip_optimizer.zero_grad()
        cross_optimizer.zero_grad()
        resnet_outputs = combied_res_clas(orig_input)
        clip_loss = combied_res_clip(clip_img_embed.to(device), orig_input)
        resnet_loss = cross_loss(resnet_outputs, label)
        print("clip损失为:", clip_loss.item())
        print("resnet损失为 :",resnet_loss.item())
        total_clip_loss += clip_loss.item()
        total_resnet_loss += resnet_loss.item()
        train_loss = 0.1*clip_loss + 0.9*resnet_loss
        train_loss.backward()
        clip_optimizer.step()
        cross_optimizer.step()   # cross_opt 包括 clip_opt优化器
        cross_scheduler.step()   # 更新 学习率

        # if (epoch%3 == 0):
        #     clip_optimizer.zero_grad()
        #     cross_optimizer.zero_grad()  # 这里进行了 验证了训练准确率 需要设置梯度为0 
        #     resnet_outputs = combied_res_clas(orig_input)

        #     clip_loss = combied_res_clip(clip_img_embed.to(device), orig_input)
        #     print("clip损失为:", clip_loss.item())
        #     clip_loss.backward()
        #     clip_optimizer.step()
        #     total_clip_loss += clip_loss.item()

        # else:

        #     # 清零梯度
        #     cross_optimizer.zero_grad()
        #     resnet_outputs = combied_res_clas(orig_input)
        #     resnet_loss = cross_loss(resnet_outputs, label)
        #     print("resnet损失为 :",resnet_loss.item())
        #     # 反向传播
        #     resnet_loss.backward()
        #     # 更新权重
        #     cross_optimizer.step()
        #     total_resnet_loss += resnet_loss.item()

        correct += (resnet_outputs.argmax(1) == label).sum().item()

        if epoch%10 == 0:
            # 创建文件夹
            os.makedirs("./model/save/cifar100", exist_ok=True)
            torch.save(combied_res_clas.state_dict(), f"./model/save/cifar100/resnet_{epoch}.pth")
        
    return total_clip_loss/len(cifar100_dataloader), total_resnet_loss/len(cifar100_dataloader), correct/len(cifar100_dataloader.dataset)



for epoch_index in range(num_epochs):
    start_t4 = time.time()
    train_clip_loss, trian_resnet_loss, train_accuracy = train(epoch_index)
    train_clip_losses.append(train_clip_loss)    
    trian_resnet_losses.append(trian_resnet_loss)    
    train_accuracies.append(train_accuracy)      
    test_accuracy = evaluate(cifar100_dataloader_test, combied_res_clas)
    test_accuracies.append(test_accuracy)
    print(f"Epoch{epoch_index+1}/{num_epochs}, Clip_Loss: {train_clip_loss:.4f}, Resnet_Loss: {trian_resnet_loss:.4f},",
                f"Train Acc: {train_accuracy:.4f}, Test Acc: {test_accuracy: .4f}")
    print("训练一个epoch耗时: ", (time.time()-start_t4)/60)


# 1 绘制损失和准确率曲线
plt.figure(figsize=(12, 5))
# 绘制损失曲线
plt.plot(train_clip_losses, label="Train CLIP Loss", color='blue')
plt.plot(trian_resnet_losses, label="Train RESNET Loss", color='green')
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training Loss Curve")
plt.legend()
plt.savefig('./model/img/training_loss_curve.png')  # 将图片保存下来

plt.figure(figsize=(12, 5))
# 2 绘制训练准确率图像
plt.plot(train_accuracies, label="Train Accuracies", color='red')
plt.plot(test_accuracies, label="Test Accuracies", color='yellow')
plt.xlabel("Epoch")
plt.ylabel("Accuracies")
plt.title("Accuracies Curve")
plt.legend()
plt.savefig('./model/img/training_and_test_accuracies_curve.png')

# 4 保存图像到文件
plt.tight_layout()
plt.close()

# 保存一些指标
save_txt(train_clip_losses, "./model/log/train_clip_losses.txt")
save_txt(trian_resnet_losses, "./model/log/trian_resnet_losses.txt")
save_txt(train_accuracies, "./model/log/train_accuracies.txt")
save_txt(test_accuracies, "./model/log/test_accuracies.txt")
